package p.ithorns.sample.influxdb.core;

import com.google.gson.*;
import com.google.gson.reflect.TypeToken;
import com.influxdb.Cancellable;
import com.influxdb.client.domain.InfluxQLQuery;
import com.influxdb.client.service.InfluxQLQueryService;
import com.influxdb.internal.AbstractQueryApi;
import com.influxdb.query.InfluxQLQueryResult;
import com.influxdb.utils.Arguments;
import okhttp3.ResponseBody;
import okio.BufferedSource;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
import org.apache.commons.csv.CSVRecord;
import org.jetbrains.annotations.NotNull;
import retrofit2.Call;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;

/**
 * InfluxQLQueryApi
 *
 * @author HtL
 * @date 2025/1/16 13:46
 * @since 1.0.0
 */
public class InfluxQLQueryApi extends AbstractQueryApi {

    private final InfluxQLQueryService service;

    public InfluxQLQueryApi(final InfluxQLQueryService service) {
        Arguments.checkNotNull(service, "service");
        this.service = service;
    }

    @NotNull
    public List<QueryRecord> query(@NotNull InfluxQLQuery query, Class<?> clazz) {
        return query(query, query.getAcceptHeader(), clazz);
    }

    @NotNull
    public List<QueryRecord> queryCsv(@NotNull InfluxQLQuery query, Class<?> clazz) {
        return query(query, InfluxQLQuery.AcceptHeader.CSV, clazz);
    }

    @NotNull
    public List<QueryRecord> queryJson(@NotNull InfluxQLQuery query, Class<?> clazz) {
        return query(query, InfluxQLQuery.AcceptHeader.JSON, clazz);
    }

    @Nonnull
    private List<QueryRecord> query(
            @Nonnull final InfluxQLQuery influxQlQuery,
            @javax.annotation.Nullable final InfluxQLQuery.AcceptHeader accept,
            Class<?> clazz
    ) {
        Call<ResponseBody> call = service.query(
                influxQlQuery.getCommand(),
                influxQlQuery.getDatabase(),
                influxQlQuery.getRetentionPolicy(),
                influxQlQuery.getPrecision() != null ? influxQlQuery.getPrecision().getSymbol() : null,
                null,
                accept != null ? accept.getVal() : InfluxQLQuery.AcceptHeader.JSON.getVal());
        AtomicReference<List<QueryRecord>> atomicReference = new AtomicReference<>();
        BiConsumer<Cancellable, BufferedSource> consumer = (cancellable, bufferedSource) -> {
            try {
                List<QueryRecord> result = parseResponse(clazz, bufferedSource, cancellable, accept);
                atomicReference.set(result);
            } catch (IOException e) {
                ERROR_CONSUMER.accept(e);
            }
        };
        query(call, consumer, ERROR_CONSUMER, EMPTY_ACTION, false);
        return atomicReference.get();
    }

    private List<QueryRecord> parseResponse(
            Class<?> clazz,
            @Nonnull final BufferedSource bufferedSource,
            @Nonnull final Cancellable cancellable,
            @Nonnull final InfluxQLQuery.AcceptHeader accept) throws IOException {

        Arguments.checkNotNull(bufferedSource, "bufferedSource");

        try (Reader reader = new InputStreamReader(bufferedSource.inputStream(), StandardCharsets.UTF_8)) {
            if (accept == InfluxQLQuery.AcceptHeader.CSV) {
                return readInfluxQLCSVResult(clazz, reader, cancellable);
            }
            return readInfluxQLJsonResult(reader, cancellable);
        }
    }

    private List<QueryRecord> readInfluxQLCSVResult(
            Class<?> clazz, @Nonnull final Reader reader, @Nonnull final Cancellable cancellable
    ) throws IOException {
        List<QueryRecord> records = new ArrayList<>();
        Map<List<Object>, QueryRecord.Series> seriesMap = null;
        Map<String, Integer> headerCols = null;
        final int nameCol = 0;
        final int tagsCol = 1;
        // The first 2 columns are static (`name`, `tags`) and got skipped.
        // All other columns are dynamically parsed
        final int dynamicColumnsStartIndex = 2;
        try (CSVParser parser = new CSVParser(reader, CSVFormat.DEFAULT.builder()
                .setIgnoreEmptyLines(false).build())) {
            for (CSVRecord csvRecord : parser) {
                if (cancellable.isCancelled()) {
                    break;
                }
                if (csvRecord.size() == 1 && csvRecord.get(0).equals("")) {
                    continue;
                }
                QueryRecord record = new QueryRecord();
                QueryRecord.Series series = null;
                if (seriesMap == null) {
                    List<String> header = csvRecord.toList();
                    headerCols = new LinkedHashMap<>();
                    for (int col = dynamicColumnsStartIndex; col < header.size(); col++) {
                        String colName = header.get(col);
                        headerCols.put(colName, col - dynamicColumnsStartIndex);
                    }
                    seriesMap = new LinkedHashMap<>();
                } else {
                    String name = csvRecord.get(nameCol);
                    Map<String, String> finalTags = parseTags(csvRecord.get(tagsCol));
                    Map<String, Integer> finalHeaderCols = headerCols;
                    series = seriesMap.computeIfAbsent(
                            Arrays.asList(name, finalTags),
                            n -> new QueryRecord.Series(name, finalTags, finalHeaderCols)
                    );
                    record.setSeries(series);

                    LinkedHashMap<String, Object> values = record.getValues();
                    Set<Map.Entry<String, Integer>> colEntries = headerCols.entrySet();
                    for (Map.Entry<String, Integer> entry : colEntries) {
                        values.put(entry.getKey(),
                                csvRecord.get(entry.getValue() + dynamicColumnsStartIndex));
                    }
                }

                if (null != series) {
                    records.add(record);
                }
            }
        }
        return records;
    }


    private static Map<String, String> parseTags(@Nonnull final String value) {
        final Map<String, String> tags = new HashMap<>();
        if (value.length() > 0) {
            for (String entry : value.split(",")) {
                final String[] kv = entry.split("=");
                tags.put(kv[0], kv[1]);
            }
        }
        return tags;
    }


    private List<QueryRecord> readInfluxQLJsonResult(Reader reader, Cancellable cancellable) {
        TypeToken<QueryRecord> typeToken = new TypeToken<>() {
        };
        Type type = typeToken.getType();
        Gson gson = new GsonBuilder()
                .registerTypeAdapter(type, new InfluxQLQueryApi.ResultsDeserializer(cancellable))
                .registerTypeAdapter(type, new InfluxQLQueryApi.ResultDeserializer())
                .create();
        try {
            return gson.fromJson(reader, type);
        } catch (JsonSyntaxException | JsonIOException jse) {
            ERROR_CONSUMER.accept(jse);
            return null;
        }
    }


    public static class ResultsDeserializer implements JsonDeserializer<List<QueryRecord>> {

        Cancellable cancellable;

        public ResultsDeserializer(final Cancellable cancellable) {
            this.cancellable = cancellable;
        }

        @Override
        public List<QueryRecord> deserialize(final JsonElement elem, final Type type,
                                             final JsonDeserializationContext ctx) throws JsonParseException {
            List<QueryRecord> results = new ArrayList<>();
            JsonObject result = elem.getAsJsonObject();
            if (result.has("results")) {
                JsonArray jsonArray = result.get("results").getAsJsonArray();
                for (JsonElement jsonElement : jsonArray) {
                    if (cancellable.isCancelled()) {
                        break;
                    }
                    results.add(ctx.deserialize(jsonElement, InfluxQLQueryResult.Result.class));
                }
            }
            return results;
        }
    }

    public static class ResultDeserializer implements JsonDeserializer<List<QueryRecord>> {
        @Override
        public List<QueryRecord> deserialize(final JsonElement elem, final Type type,
                                       final JsonDeserializationContext ctx) throws JsonParseException {
            JsonObject results = ((JsonObject) elem).getAsJsonArray("results")
                    .get(0).getAsJsonObject();
            // int id = results.get("statement_id").getAsInt();
            List<QueryRecord> records = new ArrayList<>();
            JsonArray seriesArray = results.getAsJsonArray("series");
            if (seriesArray != null) {
                for (JsonElement jsonSeries : seriesArray) {
                    JsonObject jsonObj = jsonSeries.getAsJsonObject();
                    String name = jsonObj.getAsJsonObject().get("name").getAsString();
                    Map<String, Integer> columns = new LinkedHashMap<>();
                    Map<String, String> tags = null;
                    // Handle columns
                    JsonArray jac = jsonObj.get("columns").getAsJsonArray();
                    final AtomicInteger count = new AtomicInteger(0);
                    jac.forEach(e -> {
                        columns.put(e.getAsString(), count.getAndIncrement());
                    });

                    QueryRecord.Series series = null;
                    // Handle tags - if they exist
                    if (jsonObj.get("tags") != null) {
                        JsonObject tagsObj = jsonObj.get("tags").getAsJsonObject();
                        tags = new LinkedHashMap<>();
                        for (String key : tagsObj.keySet()) {
                            tags.put(key, tagsObj.get(key).getAsString());
                        }
                        series = new QueryRecord.Series(name, tags, columns);
                    } else {
                        series = new QueryRecord.Series(name, columns);
                    }

                    JsonArray jsonArray = jsonObj.get("values").getAsJsonArray();
                    if (jsonArray != null) {
                        for (JsonElement jVal : jsonArray) {
                            QueryRecord record = new QueryRecord();
                            record.setSeries(series);

                            JsonArray jae = jVal.getAsJsonArray();
                            LinkedHashMap<String, Object> values = record.getValues();
                            Set<Map.Entry<String, Integer>> colEntries = columns.entrySet();
                            for (Map.Entry<String, Integer> entry : colEntries) {
                                values.put(entry.getKey(), jae.get(entry.getValue()).getAsString());
                            }
                            records.add(record);
                        }
                    }
                }
            }
            return records;
        }
    }

}